#!/usr/bin/env python
# coding=utf-8
# __author__ = 'Yunchao Ling'

from pymongo import MongoClient
import pymongo
import datetime
import os
import sys
import json
import time
# import requests
import StringIO
import pycurl
from retrying import retry

HOST = 'localhost'
PORT = 27017
DB_NAME = 'pmc'
COLLECTION_NAME = 'metalist'

BACK_DAY = 7

NOW = datetime.datetime.now()
NOW_DATE = NOW.strftime("%Y%m%d")


def GetConfig():
    if os.path.exists("config.ini"):
        import ConfigParser
        global HOST
        global PORT
        global DB_NAME
        global COLLECTION_NAME
        config = ConfigParser.ConfigParser()
        config.read("config.ini")
        HOST = config.get("mongodb", "host")
        PORT = int(config.get("mongodb", "port"))
        DB_NAME = config.get("mongodb", "db")
        COLLECTION_NAME = config.get("mongodb", "collection")


@retry
def Get_PMCID_List(nameString, dateStringStart, dateStringEnd):
    url = 'https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esearch.fcgi'
    buf = StringIO.StringIO()
    c = pycurl.Curl()
    c.setopt(pycurl.POST, 1)
    c.setopt(c.URL, url)
    c.setopt(pycurl.HTTPPOST,
             [('db', 'pmc'), ('mindate', dateStringStart), ('maxdate', dateStringEnd), ('retmode', "json"),
              ('retmax', "100000")])
    c.setopt(c.WRITEFUNCTION, buf.write)
    c.perform()
    result = buf.getvalue()
    print result
    print "==================================================="
    json_object = json.loads(result)
    outfile = open("pmcid_list_" + nameString + ".json", "w")
    outfile.write(buf.getvalue())
    outfile.close()
    buf.close()
    print "pmcid_list_" + nameString + " has been successfully downloaded."
    return json_object

# @retry
# def Get_PMCID_List(nameString, dateStringStart, dateStringEnd):
#     url = 'https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esearch.fcgi'
#     payload = {"db": "pmc", "mindate": dateStringStart, "maxdate": dateStringEnd, "retmode": "json",
#                "retmax": "1000000"}
#     r = requests.get(url, params=payload)
#     json_object = r.json()
#     with open("pmcid_list_" + nameString + ".json", 'wb') as fd:
#         for chunk in r.iter_content(1024 * 1024):
#             fd.write(chunk)
#     print "pmcid_list_" + nameString + " has been successfully downloaded."
#     return json_object


if __name__ == '__main__':
    print("Loading date, today is ... "),
    time.sleep(1)
    print(NOW_DATE + ".")
    time.sleep(1)

    print("Checking folder ...")
    time.sleep(1)
    if os.path.exists("pmcmeta_" + NOW_DATE):
        print("Folder pmcmeta_" + NOW_DATE + " exists, program abort.")
        sys.exit(0)
    print("Passed. Creating folder pmcmeta_" + NOW_DATE + " ... "),
    os.mkdir("pmcmeta_" + NOW_DATE)
    os.chdir("pmcmeta_" + NOW_DATE)
    time.sleep(1)
    print("complete.")
    time.sleep(1)

    print("Loading config file ... "),
    time.sleep(1)
    GetConfig()
    print("complete.")
    time.sleep(1)

    print("Initiating mongodb connection ... "),
    time.sleep(1)
    connection = MongoClient(HOST, PORT)
    db = connection[DB_NAME]
    collection = db[COLLECTION_NAME]
    print("complete.")
    time.sleep(1)

    print("Loading if PMC last update date is given ... "),
    time.sleep(1)
    if len(sys.argv) >= 2:
        LAST_UPDATE_DATE = datetime.datetime.strptime(sys.argv[1], "%Y%m%d")
        BACK_DAY = (NOW - LAST_UPDATE_DATE).days + 3
        print("yes, given date " + sys.argv[1] + ", program will go back " + str(BACK_DAY) + " days.")
    else:
        print("no.")
        print("Searching MongoDB for last update date ... "),
        results = collection.find().sort("pmc_date", pymongo.DESCENDING)
        for result in results:
            LAST_UPDATE_DATE = result["pmc_date"]
            BACK_DAY = (NOW - LAST_UPDATE_DATE).days + 3
            lastUpdateDateString = LAST_UPDATE_DATE.strftime("%Y%m%d")
            print("OK, last update date " + lastUpdateDateString + ", program will go back " + str(BACK_DAY) + " days.")
            break
    time.sleep(1)

    count_total = 0
    count_todo_new = 0
    count_todo_update = 0

    outfile_new = open("pmcid_todo.new." + NOW_DATE + ".txt", "w")
    outfile_update = open("pmcid_todo.update." + NOW_DATE + ".txt", "w")

    for i in xrange(BACK_DAY, -1, -1):
        dayDate = NOW - datetime.timedelta(days=i)
        pmcDate = datetime.datetime.combine(dayDate.date(), datetime.time(12))
        dayNameString = dayDate.strftime("%Y%m%d")
        dayDateString = dayDate.strftime("%Y/%m/%d")
        print("Downloading PMCID list " + dayNameString + " ... ")
        jsonObject = Get_PMCID_List(dayNameString, dayDateString, dayDateString)
        print("Analyzing PMCID list " + dayNameString + " ... "),
        day_count = 0
        for id in jsonObject["esearchresult"]["idlist"]:
            if id.isdigit():
                count_total += 1
                day_count += 1
                pmcid = int(id)
                result = collection.find_one({"pmcid": pmcid})
                if result == None:
                    count_todo_new += 1
                    collection.update_one({"pmcid": pmcid},
                                          {"$set": {"time_in": NOW, "time_last": NOW, "pmc_date": pmcDate}},
                                          upsert=True)
                    outfile_new.write(id + "\n")
                    outfile_new.flush()
                else:
                    if result.has_key("pmc_date"):
                        mongoDate = result["pmc_date"]
                        if pmcDate > mongoDate:
                            count_todo_update += 1
                            collection.update_one({"pmcid": pmcid}, {"$set": {"pmc_date": pmcDate, "time_last": NOW}})
                            outfile_update.write(id + "\t" + mongoDate.strftime("%Y%m%d%H") + "\n")
                            outfile_update.flush()
                    else:
                        count_todo_update += 1
                        collection.update_one({"pmcid": pmcid}, {"$set": {"pmc_date": pmcDate, "time_last": NOW}})
                        outfile_update.write(id + "\t" + str(None) + "\n")
                        outfile_update.flush()
        print("recognize " + str(day_count) + " PMCIDs.")
    connection.close()
    outfile_new.close()
    outfile_update.close()
    print("Recognize " + str(count_total) + " PMCIDs, " + str(
        count_todo_new) + " new PMCIDs and " + str(
        count_todo_update) + " update PMCIDs of meta information to be downloaded.")
    time.sleep(1)
    print "Pipeline complete."
